package com.os.repository;

import cn.hutool.json.JSONUtil;
import lombok.extern.slf4j.Slf4j;
import org.opensearch.action.search.ClearScrollRequest;
import org.opensearch.action.search.SearchRequest;
import org.opensearch.action.search.SearchResponse;
import org.opensearch.action.search.SearchScrollRequest;
import org.opensearch.client.RequestOptions;
import org.opensearch.client.RestHighLevelClient;
import org.opensearch.client.core.CountRequest;
import org.opensearch.client.core.CountResponse;
import org.opensearch.search.SearchHit;
import org.opensearch.search.SearchHits;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.List;

/**
 * 通用操作os
 */
@Slf4j
@Component
public class GeneralRepository {

    // 滚动上下文保持5分钟
    private static final String SCROLL_TIMEOUT = "5m";

    @Resource
    RestHighLevelClient client;

    public Long count(CountRequest countRequest) {
        // 执行查询
        CountResponse count = null;
        try {
            count = client.count(countRequest, RequestOptions.DEFAULT);
        } catch (Exception e) {
            log.error("openSearch count error: {}", e.getMessage());
            throw new RuntimeException(e);
        }
        // 获取查询结果
        return count.getCount();
    }

    /**
     * 根据查询条件分页
     * 数据量在10000以下的可走此方法
     * @param searchRequest
     * @return
     */
    public List<String> page(SearchRequest searchRequest) {
        return this.page(searchRequest, String.class);
    }

    /**
     * 根据查询条件分页，可指定返回类型
     * 数据量在10000以下的可走此方法
     * @param searchRequest
     * @param returnClass
     * @return
     * @param <R>
     */
    public <R> List<R> page(SearchRequest searchRequest, Class<R> returnClass) {
        List<R> productList = new ArrayList<>();
        // 执行查询
        SearchResponse searchResponse = null;
        try {
            searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
            // 获取查询结果
            SearchHits hits = searchResponse.getHits();

            hits.forEach(item -> {
                if (String.class == returnClass) {
                    productList.add((R) item.getSourceAsString());
                } else {
                    productList.add(JSONUtil.toBean(item.getSourceAsString(), returnClass));
                }
            });
        } catch (Exception e) {
            log.error("openSearch page error: {}", e.getMessage());
            throw new RuntimeException(e);
        }
        return productList;
    }

    /**
     * 根据查询条件搜索
     * @param searchRequest
     * @return
     */
    public List<String> search(SearchRequest searchRequest, Integer pageNo, Integer pageSize) {
        return this.search(searchRequest, pageNo, pageSize, String.class);
    }

    /**
     * 根据查询条件搜索，可指定返回类型
     * 默认获取一批数据（目前为1000），减少调用，然后内存比较，效率高
     * @param searchRequest
     * @param returnClass
     * @return
     * @param <R>
     */
    public <R> List<R> search(SearchRequest searchRequest, Integer pageNo, Integer pageSize, Class<R> returnClass) {
        List<R> reList = new ArrayList<>();
        String scrollId = null;
        try {
            // 执行初始查询，获取第一页数据和scrollId
            SearchResponse searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
            scrollId = searchResponse.getScrollId();
            SearchHit[] hits = searchResponse.getHits().getHits();

            // 跳过前面的页数，只返回当前页的数据
            int skip = (pageNo - 1) * pageSize;
            int currentPageHitsCount = 0;

            // 遍历所有结果，直到找到当前页的数据或遍历完所有数据
            while (hits != null && hits.length > 0) {
                for (SearchHit hit : hits) {
                    if (skip > 0) {
                        skip--; // 跳过前面的页数
                        continue;
                    }
                    if (currentPageHitsCount >= pageSize) {
                        break; // 当前页数据已满，退出循环
                    }
                    // 将命中的文档转换为re对象
                    if (String.class == returnClass) {
                        reList.add((R) hit.getSourceAsString());
                    } else {
                        reList.add(JSONUtil.toBean(hit.getSourceAsString(), returnClass));
                    }
                    currentPageHitsCount++;
                }

                // 如果当前页数据已满或没有更多数据，退出循环
                if (currentPageHitsCount >= pageSize || hits.length < pageSize) {
                    break;
                }

                // 使用scrollId获取下一页数据
                SearchScrollRequest scrollRequest = new SearchScrollRequest(scrollId);
                // scrollRequest.scroll(new Scroll(TimeValue.timeValueMinutes(5L))); // 滚动上下文保持5分钟
                scrollRequest.scroll(SCROLL_TIMEOUT);
                searchResponse = client.scroll(scrollRequest, RequestOptions.DEFAULT);
                scrollId = searchResponse.getScrollId(); // 更新scrollId
                hits = searchResponse.getHits().getHits(); // 获取下一页数据
            }
        } catch (Exception e) {
            log.error("openSearch search error: {}", e.getMessage());
            throw new RuntimeException(e);
        } finally {
            // 清理滚动上下文，释放资源
            clearScroll(scrollId);
        }
        return reList;
    }


    /**
     * 根据查询条件搜索
     * @param searchRequest
     * @return
     */
    public List<String> search(SearchRequest searchRequest) {
        return this.search(searchRequest, String.class);
    }

    /**
     * 根据查询条件搜索，可指定返回类型
     * @param searchRequest
     * @param returnClass
     * @return
     * @param <R>
     */
    public <R> List<R> search(SearchRequest searchRequest, Class<R> returnClass) {
        List<R> reList = new ArrayList<>();
        String scrollId = null;
        try {
            // 执行初始查询，获取第一批数据和scrollId
            SearchResponse searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
            scrollId = searchResponse.getScrollId();
            SearchHit[] hits = searchResponse.getHits().getHits();

            // 持续提取数据，直到所有文档都被处理完
            while (hits != null && hits.length > 0) {
                // 将当前批次的文档转换为Product对象并添加到结果列表中
                for (SearchHit hit : hits) {
                    if (String.class == returnClass) {
                        reList.add((R) hit.getSourceAsString());
                    } else {
                        reList.add(JSONUtil.toBean(hit.getSourceAsString(), returnClass));
                    }
                }
                // 使用scrollId获取下一批数据
                SearchScrollRequest scrollRequest = new SearchScrollRequest(scrollId);
                scrollRequest.scroll(SCROLL_TIMEOUT); // 滚动上下文保持5分钟
                searchResponse = client.scroll(scrollRequest, RequestOptions.DEFAULT);
                scrollId = searchResponse.getScrollId(); // 更新scrollId
                hits = searchResponse.getHits().getHits(); // 获取下一批数据
            }
        } catch (Exception e) {
            log.error("openSearch search error: {}", e.getMessage());
            throw new RuntimeException(e);
        } finally {
            // 清理滚动上下文，释放资源
            clearScroll(scrollId);
        }
        return reList;
    }

    /**
     * 清理滚动上下文，释放资源
     * @param scrollId 滚动上下文ID
     * @throws Exception 清理异常
     */
    private void clearScroll(String scrollId) {
        try {
            if (scrollId != null) {
                ClearScrollRequest clearScrollRequest = new ClearScrollRequest();
                clearScrollRequest.addScrollId(scrollId);
                client.clearScroll(clearScrollRequest, RequestOptions.DEFAULT);
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }


}
